package utils

import (
	"bytes"
	"fmt"
	"golang.org/x/image/draw"
	"golang.org/x/image/tiff"
	"image"
	"image/color"
	"image/jpeg"
	"image/png"
	"math"
	"os"
	"runtime"
	"sync"
)

type ImageChannel int

const (
	ImageChannelRed ImageChannel = iota
	ImageChannelGreen
	ImageChannelBlue
	ImageChannelFull
)

// ImageCHWByImage 获取图的CHW数据
func ImageCHWByImage(img image.Image, channel ImageChannel) ([]float32, error) {
	return getImageCHW(img, channel)
}

// ImageCHWByPath 获取图的CHW数据
func ImageCHWByPath(filePath string, channel ImageChannel) ([]float32, error) {
	img, err := ImageDecode(filePath)
	if err != nil {
		return nil, err
	}

	return getImageCHW(img, channel)
}

// ImageDecode 解析图片
func ImageDecode(filePath string) (image.Image, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer func(f *os.File) {
		err = f.Close()
		if err != nil {
			fmt.Println("system error:" + err.Error())
		}
	}(f)

	// 注册图片格式
	registerImageFormat()
	img, _, err := image.Decode(f)
	return img, err
}

// ImageDecodeToRGBA 图片转为RGBA
func ImageDecodeToRGBA(filePath string, width, height int) (*image.RGBA, error) {
	f, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer func(f *os.File) {
		err = f.Close()
		if err != nil {
			fmt.Println("system error:" + err.Error())
		}
	}(f)

	// 注册图片格式
	registerImageFormat()
	img, _, err := image.Decode(f)
	if err != nil {
		return nil, err
	}

	if width > 0 && height > 0 {
		return ImageResizeToRGBA(img, width, height), nil
	} else {
		return ImageResizeToRGBA(img, img.Bounds().Dx(), img.Bounds().Dy()), nil
	}
}

// ImageToGrayLevel 图像转换为灰度图
func ImageToGrayLevel(img image.Image) image.Image {
	grayImg := image.NewGray(img.Bounds())
	for y := img.Bounds().Min.Y; y < img.Bounds().Max.Y; y++ {
		for x := img.Bounds().Min.X; x < img.Bounds().Max.X; x++ {
			// 获取彩色图像的像素值
			colorRGB := img.At(x, y)
			r, g, b, _ := colorRGB.RGBA()

			// 计算灰度值并设置到灰度图像中
			grayValue := uint8(0.299*float64(r) + 0.587*float64(g) + 0.114*float64(b))
			gray := color.Gray{Y: grayValue}
			grayImg.Set(x, y, gray)
		}
	}

	return grayImg
}

// ImageRawDataToGrayPng 图片数据转png灰度图
func ImageRawDataToGrayPng(rawData []float32, width, height int) (bytes.Buffer, error) {
	oImage := image.NewGray(image.Rect(0, 0, width, height))
	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			oImage.Set(i, j, color.Gray{Y: uint8(rawData[i+width*j] * 255)})
		}
	}
	var imageBuf bytes.Buffer
	err := png.Encode(&imageBuf, oImage)
	if err != nil {
		return bytes.Buffer{}, err
	}

	return imageBuf, nil
}

// ImageRawPercentDataToGrayPng 图片数据转png灰度图
func ImageRawPercentDataToGrayPng(rawData []float32, width, height int) (bytes.Buffer, error) {

	oImage := image.NewGray(image.Rect(0, 0, width, height))
	for i := 0; i < width; i++ {
		for j := 0; j < height; j++ {
			if rawData[i+width*j] > 0.5 {
				oImage.Set(i, j, color.Gray{Y: uint8(255)})
			} else {
				oImage.Set(i, j, color.Gray{Y: uint8(0)})
			}

		}
	}

	var imageBuf bytes.Buffer
	err := png.Encode(&imageBuf, oImage)
	if err != nil {
		return bytes.Buffer{}, err
	}

	return imageBuf, nil
}

// ImageResizeToRGBA 图片缩放
func ImageResizeToRGBA(img image.Image, width, height int) *image.RGBA {
	resizedImg := image.NewRGBA(image.Rect(0, 0, width, height))
	draw.BiLinear.Scale(resizedImg, resizedImg.Bounds(), img, img.Bounds(), draw.Over, nil)
	return resizedImg
}

// ImageResize 图片缩放
func ImageResize(img image.Image, width, height int) image.Image {
	resizedImg := image.NewRGBA(image.Rect(0, 0, width, height))
	draw.BiLinear.Scale(resizedImg, resizedImg.Bounds(), img, img.Bounds(), draw.Over, nil)
	return resizedImg
}

// RGBAResize 图片缩放
func RGBAResize(img *image.RGBA, width, height int) *image.RGBA {
	resizedImg := image.NewRGBA(image.Rect(0, 0, width, height))
	draw.BiLinear.Scale(resizedImg, resizedImg.Bounds(), img, img.Bounds(), draw.Over, nil)
	return resizedImg
}

// AiImgOutput 处理ai图像把白色变为绿色和原图叠加输出
func AiImgOutput(inputImg image.Image, output []float32) (image.Image, error) {

	resizedImg := image.NewRGBA(inputImg.Bounds())
	draw.Draw(resizedImg, inputImg.Bounds(), inputImg, inputImg.Bounds().Min, draw.Src)
	// 处理输出并修改图像
	for y := 0; y < resizedImg.Bounds().Dy(); y++ {
		for x := 0; x < resizedImg.Bounds().Dx(); x++ {
			if output[y*resizedImg.Bounds().Dx()+x] == 1 {
				r, g, b, _ := resizedImg.At(x, y).RGBA()
				resizedImg.Set(x, y, color.RGBA{
					R: uint8(float32(r) * 0.7 / 0x101),
					G: uint8(float32(g) / 0x101),
					B: uint8(float32(b) * 0.7 / 0x101),
					A: 255,
				})
			}
		}
	}

	return resizedImg, nil
}

// AiImgOutputResNet 处理ai图像把白色变为绿色和原图叠加输出
func AiImgOutputResNet(inputImg *image.RGBA, output []float32) (image.Image, error) {
	// 处理输出并修改图像
	for y := 0; y < inputImg.Bounds().Dy(); y++ {
		for x := 0; x < inputImg.Bounds().Dx(); x++ {
			if output[y*inputImg.Bounds().Dx()+x] > 0.5 {
				r, g, b, _ := inputImg.At(x, y).RGBA()
				inputImg.Set(x, y, color.RGBA{
					R: uint8(float32(r) * 0.7 / 0x101),
					G: uint8(float32(g) / 0x101),
					B: uint8(float32(b) * 0.7 / 0x101),
					A: 255,
				})
			}
		}
	}

	return inputImg, nil
}

// ImageToRgbaConcurrence 图片转rgba并并发处理
func ImageToRgbaConcurrence(img image.Image) *image.RGBA {
	bounds := img.Bounds()
	rgba := image.NewRGBA(bounds)

	var wg sync.WaitGroup
	pixelCh := make(chan struct{ x, y int })

	for i := 0; i < runtime.NumCPU(); i++ {
		go func() {
			for pixel := range pixelCh {
				rgba.Set(pixel.x, pixel.y, img.At(pixel.x, pixel.y))
				wg.Done()
			}
		}()
	}

	for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
		for x := bounds.Min.X; x < bounds.Max.X; x++ {
			wg.Add(1)
			pixelCh <- struct{ x, y int }{x, y}
		}
	}

	wg.Wait()
	close(pixelCh)

	return rgba
}

func registerImageFormat() {
	image.RegisterFormat("png", "png", png.Decode, png.DecodeConfig)
	image.RegisterFormat("jpg", "jpg", jpeg.Decode, jpeg.DecodeConfig)
	image.RegisterFormat("tiff", "tiff", tiff.Decode, tiff.DecodeConfig)
}

// 获取像素值
func getImageCHW(img image.Image, imageChannel ImageChannel) ([]float32, error) {
	bounds := img.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y

	outLength := width * height
	if imageChannel == ImageChannelFull {
		outLength = width * height * 3
	}
	chwData := make([]float32, outLength)

	idx := 0
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			c := img.At(x, y)
			r, g, b, _ := c.RGBA()

			switch imageChannel {
			case ImageChannelFull:
				chwData[idx] = float32(r) / 65535.0
				chwData[idx+height*width] = float32(g) / 65535.0
				chwData[idx+2*height*width] = float32(b) / 65535.0
				break
			case ImageChannelRed:
				chwData[idx] = float32(math.Floor((float64(r)/65535.0)*255.0)) / 255.0
				break
			case ImageChannelBlue:
				chwData[idx] = float32(math.Floor((float64(b)/65535.0)*255.0)) / 255.0
				break
			case ImageChannelGreen:
				chwData[idx] = float32(math.Floor((float64(g)/65535.0)*255.0)) / 255.0
				break
			}

			idx++
		}
	}

	return chwData, nil
}
